package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"strings"

	"gopkg.in/yaml.v3"
)

// UDPRule 定义UDP转发规则
type UDPRule struct {
	LocalPort  int    `json:"local_port" yaml:"local_port"`
	TargetAddr string `json:"target_addr" yaml:"target_addr"`
}

// UDPConfig 定义UDP代理配置
type UDPConfig struct {
	Rules []UDPRule `json:"rules" yaml:"rules"`
}

func loadConfigFromFile(filePath string) (UDPConfig, error) {
	var config UDPConfig
	data, err := os.ReadFile(filePath)
	if err != nil {
		return config, err
	}
	if strings.HasSuffix(filePath, ".json") {
		err = json.Unmarshal(data, &config)
	} else if strings.HasSuffix(filePath, ".yaml") || strings.HasSuffix(filePath, ".yml") {
		err = yaml.Unmarshal(data, &config)
	} else {
		return config, fmt.Errorf("unsupported config file format")
	}
	if err != nil {
		return config, err
	}
	return config, nil
}

func loadConfigFromEnv() []UDPRule {
	var rules []UDPRule
	envRules := os.Getenv("UDP_PROXY_RULES")
	if envRules == "" {
		return rules
	}
	for _, ruleStr := range strings.Split(envRules, ";") {
		parts := strings.SplitN(ruleStr, ":", 3)
		if len(parts) != 3 {
			log.Printf("无效的环境变量规则格式: %s\n", ruleStr)
			continue
		}
		localPort, err := strconv.Atoi(parts[1])
		if err != nil {
			log.Printf("无法解析本地端口: %s\n", parts[1])
			continue
		}
		rule := UDPRule{
			LocalPort:  localPort,
			TargetAddr: parts[2],
		}
		rules = append(rules, rule)
	}
	return rules
}

func forwardUDP(localConn *net.UDPConn, targetAddr *net.UDPAddr) {
	buffer := make([]byte, 1500) // 以太网MTU为1500字节，优化性能
	for {
		n, _, err := localConn.ReadFromUDP(buffer)
		if err != nil {
			log.Printf("从本地读取数据失败: %v\n", err)
			continue
		}
		targetConn, err := net.DialUDP("udp", nil, targetAddr)
		if err != nil {
			log.Printf("连接目标地址 %s 失败: %v\n", targetAddr.String(), err)
			continue
		}
		_, err = targetConn.Write(buffer[:n])
		targetConn.Close()
		if err != nil {
			log.Printf("向目标 %s 写入数据失败: %v\n", targetAddr.String(), err)
		}
	}
}

func reverseForwardUDP(targetConn *net.UDPConn, localAddr *net.UDPAddr) {
	buffer := make([]byte, 1500)
	for {
		n, err := targetConn.Read(buffer)
		if err != nil {
			log.Printf("从目标读取数据失败: %v\n", err)
			continue
		}
		conn, err := net.DialUDP("udp", nil, localAddr)
		if err != nil {
			log.Printf("连接本地地址 %s 失败: %v\n", localAddr.String(), err)
			continue
		}
		_, err = conn.Write(buffer[:n])
		conn.Close()
		if err != nil {
			log.Printf("向本地写入数据失败: %v\n", err)
		}
	}
}

func deduplicateRules(rules []UDPRule) []UDPRule {
	ruleMap := make(map[string]bool)
	var uniqueRules []UDPRule
	for _, rule := range rules {
		key := fmt.Sprintf("%d:%s", rule.LocalPort, rule.TargetAddr)
		if _, exists := ruleMap[key]; !exists {
			ruleMap[key] = true
			uniqueRules = append(uniqueRules, rule)
		}
	}
	return uniqueRules
}

func main() {

	// 检查并创建logs目录
	if _, err := os.Stat("./logs"); os.IsNotExist(err) {
		if err := os.MkdirAll("./logs", 0755); err != nil {
			log.Fatalf("无法创建logs目录: %v", err)
		}
	}
	// 打开日志文件
	logFile, err := os.OpenFile("./logs/udp_proxy.log", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("无法打开日志文件: %v", err)
	}

	defer logFile.Close()

	multiWriter := io.MultiWriter(logFile, os.Stdout)

	// 设置日志输出为MultiWriter
	log.SetOutput(multiWriter)

	configFile := "udp_proxy.yaml"
	config, err := loadConfigFromFile(configFile)
	if err != nil {
		log.Printf("无法从文件 %s 加载配置: %v\n", configFile, err)
		return
	}

	envRules := loadConfigFromEnv()
	config.Rules = append(config.Rules, envRules...)
	config.Rules = deduplicateRules(config.Rules)

	for _, rule := range config.Rules {
		localAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", rule.LocalPort))
		if err != nil {
			log.Printf("解析本地地址失败: %v\n", err)
			continue
		}

		targetAddr, err := net.ResolveUDPAddr("udp", rule.TargetAddr)
		if err != nil {
			log.Printf("解析目标地址失败: %v\n", err)
			continue
		}

		localConn, err := net.ListenUDP("udp", localAddr)
		if err != nil {
			log.Printf("监听本地端口 %d 失败: %v\n", rule.LocalPort, err)
			continue
		}
		defer localConn.Close()

		targetConn, err := net.DialUDP("udp", nil, targetAddr)
		if err != nil {
			log.Printf("连接目标地址 %s 失败: %v\n", rule.TargetAddr, err)
			continue
		}
		defer targetConn.Close()

		log.Printf("UDP代理已启动，监听本地端口 %d 并转发到 %s\n", rule.LocalPort, rule.TargetAddr)

		go forwardUDP(localConn, targetAddr)
		go reverseForwardUDP(targetConn, localAddr)
	}

	select {}
}
